package org.eoti.reflect;

import com.sun.tools.attach.VirtualMachine;
import com.sun.tools.attach.AttachNotSupportedException;
import com.sun.tools.attach.AgentInitializationException;
import com.sun.tools.attach.AgentLoadException;

import javax.management.MBeanServerConnection;
import javax.management.ObjectName;
import javax.management.MalformedObjectNameException;
import javax.management.remote.JMXConnector;
import javax.management.remote.JMXConnectorFactory;
import javax.management.remote.JMXServiceURL;
import java.io.BufferedReader;
import java.io.File;
import java.io.IOException;
import java.io.InputStreamReader;
import java.lang.management.ManagementFactory;
import java.lang.management.ThreadInfo;
import java.lang.management.ThreadMXBean;
import java.util.*;
import java.util.concurrent.CopyOnWriteArrayList;

/**
 * The MIT License
 * <p/>
 * Copyright (c) 2009 Malachi de AElfweald (malachid@gmail.com)
 * <p/>
 * Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated
 * documentation files (the "Software"), to deal in the Software without restriction, including without limitation the
 * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to
 * permit persons to whom the Software is furnished to do so, subject to the following conditions:
 * <p/>
 * The above copyright notice and this permission notice shall be included in all copies or substantial portions of the
 * Software.
 * <p/>
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE
 * WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
 * COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
 * OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 */

// Derived from http://java.sun.com/mailers/techtips/corejava/2007/tt0807.html#1
public class ThreadUtil
implements Iterable<ThreadDetail>
{
	private CopyOnWriteArrayList<ThreadDetail> details;

	public static List<String> getAllJavaProcessIDs()
			throws IOException, InterruptedException
	{
		ArrayList<String> ids = new ArrayList<String>();
		Process p = new ProcessBuilder("jps", "-q").start();
		BufferedReader in = new BufferedReader(new InputStreamReader(p.getInputStream()));
		String line = null;
		while( (line = in.readLine()) != null)
			ids.add(line.trim());

		p.waitFor();
		return ids;
	}

	public static void main(String[] args)
	{
		try{
			List<String> ids = (args.length == 0 ? getAllJavaProcessIDs() : Arrays.asList(args));
			for(String id : ids)
			{
				try{
					ThreadUtil tu = new ThreadUtil(id);
					System.out.format("Threads for process #%s:\n", id);
					for(ThreadDetail detail : tu)
						System.out.format("\t%s\n", detail);
				}catch(Exception e){
					System.err.format("ERROR reading process #%s: %s\n", id, e.getMessage());
				}
			}
		}catch(Exception e2){
			System.err.format("Exception: %s\n", e2.getMessage());
		}
	}

	public ThreadUtil(String id)
			throws IOException, AttachNotSupportedException, AgentInitializationException, AgentLoadException, MalformedObjectNameException
	{
		details = new CopyOnWriteArrayList<ThreadDetail>();
		VirtualMachine vm = VirtualMachine.attach(id);
    String connectorAddr = vm.getAgentProperties().getProperty("com.sun.management.jmxremote.localConnectorAddress");
    if(connectorAddr == null)
    {
	    String agent = String.format(
			    "%s%slib%smanagement-agent.jar",
	        vm.getSystemProperties().getProperty("java.home"),
	        File.separator,
	        File.separator
	    );
      vm.loadAgent(agent);
      connectorAddr = vm.getAgentProperties().getProperty("com.sun.management.jmxremote.localConnectorAddress");
    }
		
    JMXServiceURL serviceURL = new JMXServiceURL(connectorAddr);
    JMXConnector connector = JMXConnectorFactory.connect(serviceURL);
    MBeanServerConnection mbsc = connector.getMBeanServerConnection();
    ObjectName objName = new ObjectName(ManagementFactory.THREAD_MXBEAN_NAME);
    Set<ObjectName> mbeans = mbsc.queryNames(objName, null);
    for( ObjectName name: mbeans)
    {
      ThreadMXBean threadBean = ManagementFactory.newPlatformMXBeanProxy(
		      mbsc,
		      name.toString(),
		      ThreadMXBean.class
      );
      long threadIds[] = threadBean.getAllThreadIds();
      for(long threadId: threadIds)
      {
        ThreadInfo threadInfo = threadBean.getThreadInfo(threadId);
	      details.add(new ThreadDetail(id, threadId, threadInfo.getThreadName(), threadInfo.getThreadState()));
      }
    }
  }

	/**
	 * Returns an iterator over a set of elements of type T.
	 *
	 * @return an Iterator.
	 */
	public Iterator<ThreadDetail> iterator()
	{
		return details.iterator();
	}
}
